package com.mongo.controller;


import cn.hutool.core.date.DateUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.text.StrPool;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.IndexOptions;
import com.mongodb.client.model.Indexes;
import lombok.extern.slf4j.Slf4j;
import org.bson.Document;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

@Slf4j
@RestController
@RequestMapping("/import")
public class ImportABCompartmentController {

    //10000
    private static final int BATCH_SIZE = 5000;
    private static final String DELIMITER = "\t"; // 使用制表符作为分隔符
    @Autowired
    private MongoTemplate mongoTemplate;

    @GetMapping("abcompartment")
    public Object importIpr() throws IOException {

        String dirPath = "D:\\work\\doc\\多组学数据库\\epigenetics\\hic\\abcompartment\\split_files";
        List<String> fileNames = FileUtil.listFileNames(dirPath);
        Collections.sort(fileNames);
        Integer total = 0;
        for (String fileName : fileNames) {
            if (FileUtil.exist(dirPath + File.separator + fileName)) {
                String collectionName = StrUtil.subBefore(StrUtil.subBefore(fileName, "_", true), StrPool.DOT, true);
                // 创建position字段的普通索引
                //
                System.out.println("开始导入集合:" + collectionName);
                System.out.println("导入的文件:" + fileName);
                total += importTxtToMongoDB(dirPath + File.separator + fileName, collectionName);
                  /*  createPositionIndex(collection);
                    createChrIndex(collection);*/
                //createUniqueIndex(collection);
                System.out.println("数据导入成功！");
            }
        }
        return JSONUtil.toJsonStr("导入成功!共导入" + total + "条数据");
    }
    Integer importTxtToMongoDB(String filePath, String collectionName) throws IOException {
        MongoCollection<Document> collection = mongoTemplate.getCollection(collectionName);
        BufferedReader br = null;
        try {
            br = new BufferedReader(new FileReader(filePath));
            String headerLine = br.readLine();
            if (headerLine == null) {
                System.out.println("文件为空，没有数据可导入。");

            }
            List<String> headers = Arrays.asList(headerLine.split(DELIMITER));
            List<Document> batch = new ArrayList<>(BATCH_SIZE);
            int count = 0;
            int lineNumber = 2; // 行号从2开始，因为第一行是表头
            String line;
            String[] longHeaders = new String[]{"Start", "End", "_id"};
            String[] doubleHeaders = new String[]{"Eigv_compartment_val"};
            String[] intHeaders = new String[]{"Gene_num"};
            while ((line = br.readLine()) != null) {
                /*if (lineNumber >= 500000) {
                    System.out.println("已导入超过500000条数据,停止导入");
                    break;
                }*/
                if (StrUtil.isNotEmpty(line)) {
                    List<String> values = Arrays.asList(line.split(DELIMITER));
                    Document doc = new Document();
                    // 添加_id字段，值为行号作为主键
                    /* doc.put("_id", lineNumber);*/

                    for (int i = 0; i < headers.size(); i++) {
                        if (i < values.size()) {
                            String headerName = headers.get(i);
                            String headerValue = values.get(i);
                            if (ArrayUtil.contains(longHeaders, headerName)) {
                                doc.put(headerName, Long.valueOf(headerValue)); // 去除字段两边的空白字符
                            } else if (ArrayUtil.contains(doubleHeaders, headerName)) {
                                doc.put(headerName, Double.valueOf(headerValue));
                            } else if (ArrayUtil.contains(intHeaders, headerName)) {
                                doc.put(headerName, Integer.valueOf(headerValue));
                            } else {
                                doc.put(headerName, headerValue.trim()); // 去除字段两边的空白字符
                            }
                        } else {
                            doc.put(headers.get(i), "");
                        }
                    }
                    batch.add(doc);
                    lineNumber++;
                    count++;

                    if (count % BATCH_SIZE == 0) {
                        collection.insertMany(batch);
                        batch.clear();
                        System.out.println(DateUtil.now() + ",已导入 " + count + " 条记录");
                    }
                }
            }
            // 导入剩余的数据
            if (!batch.isEmpty()) {
                collection.insertMany(batch);
                System.out.println(DateUtil.now() + ",总共导入了 " + count + " 条记录");
            }
            return count;
        } finally {
            if (br != null) {
                try {
                    br.close();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
            //删除当前文件
            System.out.println("导入完成,准备删除文件:" + filePath);
            FileUtil.del(filePath);
        }
    }
}
